package com.yiwenup.struct._04_heap;

import com.yiwenup.printer.BinaryTreeInfo;
import com.yiwenup.struct._04_heap.base.AbstractHeap;

import java.util.Collection;
import java.util.Comparator;

/**
 * 二叉堆：底层使用数组实现
 **/
@SuppressWarnings("unchecked")
public class BinaryHeap<E> extends AbstractHeap<E> implements BinaryTreeInfo {
    private E[] elements;

    public static final int DEFAULT_CAPACITY = 10;

    public BinaryHeap(E[] elements) {
        this(elements, null);
    }

    public BinaryHeap(E[] elements, Comparator<E> comparator) {
        super(comparator);
        if (elements == null || elements.length == 0) {
            this.elements = (E[]) new Object[DEFAULT_CAPACITY];
        } else {
            size = elements.length;
            int capacity = Math.max(elements.length, DEFAULT_CAPACITY);
            this.elements = (E[]) new Object[capacity];
            for (int i = 0; i < elements.length; i++) {
                this.elements[i] = elements[i];
            }
            heapify();
        }
    }

    public BinaryHeap(Collection<E> elements, Comparator<E> comparator) {
        super(comparator);
        size = elements == null ? 0 : elements.size();
        if (size == 0) {
            this.elements = (E[]) new Object[DEFAULT_CAPACITY];
        } else {
            int capacity = Math.max(elements.size(), DEFAULT_CAPACITY);
            this.elements = (E[]) new Object[capacity];
            int i = 0;
            for (E element : elements) {
                this.elements[i++] = element;
            }
            heapify();
        }
    }

    /**
     * 批量建堆
     */
    private void heapify() {
        // 自上而下的上滤 (nlogn)
//        for (int i = 1; i < size; i++) {
//            siftUp(i);
//        }
        // 自下而上的下滤 (n)
        for (int i = (size >> 1) - 1; i >= 0; i--) {
            siftDown(i);
        }
    }

    public BinaryHeap() {
        this((E[]) null, null);
    }

    public BinaryHeap(Comparator<E> comparator) {
        this((Collection<E>) null, comparator);
    }

    @Override
    public void clear() {
        for (int i = 0; i < elements.length; i++) {
            elements[i] = null;
        }
        size = 0;
    }

    @Override
    public void add(E element) {
        elementNotNullCheck(element);
        ensureCapacity(size + 1);
        elements[size++] = element;
        siftUp(size - 1);
    }

    public void addAll(Collection<E> elements) {
        if (elements == null) return;
        for (E element : elements) {
            add(element);
        }
    }

    /**
     * 上滤
     *
     * @param index 索引位置
     */
    private void siftUp(int index) {
//        E node = elements[index];
//        while (index > 0) {
//            int pIndex = (index - 1) >> 1;
//            E parent = elements[pIndex];
//            if (compare(node, parent) <= 0) {
//                return;
//            }
//            // 交换数组中的物理位置
//            E tmp = elements[pIndex];
//            elements[pIndex] = elements[index];
//            elements[index] = tmp;
//            // 重新赋值index
//            index = pIndex;
//        }
        E node = elements[index];
        while (index > 0) {
            int pIndex = (index - 1) >> 1;
            E parent = elements[pIndex];
            if (compare(node, parent) <= 0) break;
            // 交换数组中的物理位置
            elements[index] = parent;
            // 重新赋值index
            index = pIndex;
        }
        elements[index] = node;
    }

    @Override
    public E get() {
        emptyCheck();
        return elements[0];
    }

    @Override
    public E remove() {
        emptyCheck();

        E top = elements[0];

        int lastIndex = --size;
        elements[0] = elements[lastIndex];
        elements[lastIndex] = null;

        siftDown(0);

        return top;
    }

    private void siftDown(int index) {
        E element = elements[index];
        // index < 第一个叶子节点索引（层序遍历，找到第一个叶子节点之后的节点都是叶子节点）
        // 第一个叶子节点索引 = 非叶子节点的数量
        int half = size >> 1;
        while (index < half) {

            // index 一定有左子节点
            int cIndex = (index << 1) + 1;
            E child = elements[cIndex];
            // 可能会有右子节点
            int rIndex = cIndex + 1;
            if (rIndex < size && compare(elements[rIndex], child) > 0) {
                child = elements[cIndex = rIndex];
            }

            if (compare(element, child) >= 0) break;

            elements[index] = child;
            index = cIndex;
        }
        elements[index] = element;
    }

    @Override
    public E replace(E element) {
        elementNotNullCheck(element);

        E top = null;
        if (size == 0) {
            elements[0] = element;
            size++;
        } else {
            top = elements[0];
            elements[0] = element;
            siftDown(0);
        }

        return top;
    }


    private void emptyCheck() {
        if (size == 0) {
            throw new IndexOutOfBoundsException("Heap is empty");
        }
    }

    private void elementNotNullCheck(E element) {
        if (element == null) {
            throw new IllegalArgumentException("element must not be null");
        }
    }

    /**
     * 扩容逻辑
     *
     * @param capacity 原容量
     */
    private void ensureCapacity(int capacity) {
        int oldCap = this.elements.length;
        if (oldCap >= capacity) {
            return;
        }
        int newCap = oldCap + (oldCap >> 1);
        E[] newElements = (E[]) new Object[newCap];
        for (int i = 0; i < this.size; i++) {
            newElements[i] = this.elements[i];
        }
        this.elements = newElements;

        System.out.println(oldCap + "扩容为" + newCap);
    }

    @Override
    public Object root() {
        return 0;
    }

    @Override
    public Object left(Object node) {
        int index = ((int) node << 1) + 1;
        return index >= size ? null : index;
    }

    @Override
    public Object right(Object node) {
        int index = ((int) node << 1) + 2;
        return index >= size ? null : index;
    }

    @Override
    public Object string(Object node) {
        return elements[(int) node];
    }
}
